#!/usr/bin/env python

# xml tree <--> sqltree converter

__revision__  = "$Id: xmlsql.py 7 2007-03-18 13:29:17Z fvamail $"

from xml.sax import make_parser , saxutils
from StringIO import StringIO
from xml.sax.handler import feature_namespaces
from codecs import EncodedFile

from sql import sql_quote
import table
import baseSQLtypes
#-----------------------------------------------------------

class XmlSql(saxutils.DefaultHandler) :

    def __init__ ( self , conn , schema='public' , tables=None ,
                   xml2sqltags = None , xml2sqlattrs = None ,
                   sql2xmltags = None , sql2xmlattrs = None ) :      
        self.conn = conn
        self.schema = schema
        
        self.tables = None
        if tables : self.tables = dict([(t.table_name,t) for t in tables])
        #for t in tables : print t.table_name
        self.xml2sqltags = xml2sqltags
        self.xml2sqlattrs = xml2sqlattrs
        self.sql2xmltags = sql2xmltags
        self.sql2xmlattrs = sql2xmlattrs
        
        self.cursor = self.conn.cursor()
        self.parser = make_parser()
        self.parser.setFeature(feature_namespaces, 0)
        self.parser.setContentHandler(self)
        self.document = None
        self._to_utf()

    def _to_utf ( self ) : self.cursor.execute("set client_encoding to 'utf-8'")

    def register_table ( self , table ) :
        if self.tables is None : self.tables = {}
        self.tables[table.table_name]=table
    
    # xml -> sql part
    def insert ( self , file_or_string ) :
        if isinstance(file_or_string,basestring) :
            file_or_string = StringIO(file_or_string)
        self.parser.parse(file_or_string)
        return self.document

    def startDocument ( self ) :
        self.document = self.get_new_id ( 'document' )
        self.fragment = self.get_new_id ( 'fragment' )
        self._characters = []
        self.nodestack = []
        self.counter = -1

    def _flush ( self ) :
        tmp = self._characters
        self._characters = []
        self.startElement ( "#TEXT" , {'#PCDATA':''.join(tmp)} )
        self.endElement ( "#TEXT" )

    def endDocument ( self ) :
        if self._characters : self._flush()
        self.cursor.execute("""
        insert into %s.document (document,fragment,map)
        values (%d,%d,0)""" % ( 
            self.schema , self.document , self.fragment )) 
        #self.conn.commit()
        
    def startElement ( self , name , attrs ) :
        if self._characters : self._flush()
        if self.xml2sqltags :
            name = self.xml2sqltags(name)
        if self.xml2sqlattrs :
            attrs = dict([
                (self.xml2sqlattrs(k.encode('utf-8')),v.encode('utf-8'))
                 for (k,v) in attrs.items()])
        else : attrs = dict([(k.encode('utf-8'),v.encode('utf-8'))
                             for (k,v) in attrs.items()])
        attrs[''] = name.encode('utf-8')
        attr_list = self.attributes ( attrs )
        self.element ( attr_list )
        
    def endElement ( self , name ) :
        if self._characters : self._flush()
        self.counter += 1
        flft,attr_list = self.nodestack.pop()
        self.cursor.execute("""
        insert into %s.fragment_forest (fragment,flft,frgt,attr_list)
        values (%d,%d,%d,%d)""" % (
            self.schema , self.fragment , flft , self.counter , attr_list )) 

    def characters ( self , content ) :
        self._characters.append ( content )

    def get_new_id ( self , prefix ) :
        self.cursor.execute ( "select nextval('%s.%s_seq')" %
                              (self.schema,prefix) )
        return self.cursor.fetchone()[0]

    def insert_attr ( self , attr_list , name , value ) :
        self.cursor.execute (
            """select %s.new_attr_pair(%d,'%s','%s');"""
            % (self.schema,attr_list,sql_quote(name),sql_quote(value)))

    def attributes ( self , attrs ) :
        attr_list = self.get_new_id( 'attr_list' )
        self.insert_list ( attrs , attr_list )
        return attr_list

    def insert_list ( self , attrs , attr_list ) :
        table = self.tables and self.tables.get(attrs[''])
        #for k,v in attrs.iteritems() : print k , ':' , v , type(k) , type(v)
        if table : self.table_insert_attrs(table,attr_list,attrs)
        if attrs : self.insert_attrs(attr_list,attrs)

    def insert_attrs ( self , attr_list , attrs ) :
        for name,value in attrs.iteritems() :
            self.insert_attr( attr_list , name , value )

    def table_insert_attrs ( self , table , attr_list , attrs ) :
        assert 'attr_list' not in attrs
        attrs['attr_list'] = attr_list
        del attrs['']
        checked = table.valid(attrs)

##         print 
##         for k,v in attrs.iteritems() :
##             print 'a' , k , ':' , v , type(k) , type(v)
##         for k,v in checked.iteritems() :
##             print 'c' , k , ':' , v , type(k) , type(v)

        for k in checked : attrs.pop(k,0)
        assert not attrs or not table.strict
        row = table(checked)
        row.save(self.cursor)

    def element ( self , attr_list ) :
        self.counter += 1
        self.nodestack.append ((self.counter,attr_list))
    # end of xml -> sql part

    def fill ( self , table_name ) :
        assert table_name in self.tables
        schema = self.schema
        cr = self.conn.cursor()
        crdel = self.conn.cursor()
        cr.execute ( """
        select attr_list , name , value  from
        (
        select distinct attr_list from
        %(schema)s.attr_list
        join %(schema)s.attr_pair using (attr_pair)
        join %(schema)s.name using (name_id)
        join %(schema)s.value using (value_id)
        where name = '' and value = '%(table_name)s'
        ) as attr_lists 
        join %(schema)s.attr_list using (attr_list)
        join %(schema)s.attr_pair using (attr_pair)
        join %(schema)s.name using (name_id)
        join %(schema)s.value using (value_id)
        order by attr_list
        """ % vars() )

        row = cr.fetchone()
        if not row : return
        attr_list , name , value = row
        while 1 :
            attrs = {name:value}
            old_attr_list = attr_list
            while 1 :
                row = cr.fetchone()
                if not row : break
                attr_list , name , value = row
                if old_attr_list != attr_list : break
                attrs[name]=value
##             print '--------' , old_attr_list
##             for k,v in attrs.iteritems() :
##                 print k , ':' , v , type(k) , type(v)
            crdel.execute ( """
            delete from %(schema)s.attr_list
            where attr_list=%(old_attr_list)d
            """ % vars() )
            self.insert_list ( attrs , old_attr_list )
            if not row : break

    def free ( self , table_name ) :
        assert table_name in self.tables
        t = self.tables.pop(table_name)
        if not self.tables : self.tables = None
        cn = self._case_name ( t )
        cv = self._case_value ( t )
        cw = self._where ( t )
        ss = self._subselect ( t )
        table_schema = t.schema
        cr = self.conn.cursor()
        crdel = self.conn.cursor()
        cr.execute ( """
        select attr_list , %(cn)s as name , %(cv)s as value 
        from %(table_schema)s.%(table_name)s , %(ss)s
        where %(cw)s
        order by attr_list
         """ % vars() )

        row = cr.fetchone()
        if not row : return
        attr_list , name , value = row
        while 1 :
            attrs = {name:value}
            old_attr_list = attr_list
            while 1 :
                row = cr.fetchone()
                if not row : break
                attr_list , name , value = row
                if old_attr_list != attr_list : break
                attrs[name]=value
##             print '--------' , old_attr_list
##             for k,v in attrs.iteritems() :
##                 print k , ':' , v , type(k) , type(v)
            crdel.execute ( """
            delete from %(table_schema)s.%(table_name)s
            where attr_list=%(old_attr_list)d
            """ % vars() )
            self.insert_list ( attrs , old_attr_list )
            if not row : break

    # sql -> xml part
    def filter ( self , documents ) :   # [(1,4),(2,2),(0,10)...]
        """Yields existed documents only (in order requested)"""
        cond = 'or'.join(
            ['(document between %d and %d)'%p for p in documents])
        cursor = self.conn.cursor()
        cursor.execute (
            'select document from %s.document where %s' % (self.schema,cond))
        exist = {}
        while 1 :
            row = cursor.fetchone()
            if not row : break
            exist[row[0]] = 1
        maxdocument = (exist and max(exist.keys()) or 0) + 1
        for b,e in documents :
            for d in range(b,min(e+1,maxdocument)) :
                if d in exist : yield d

    def _stack_update ( self , writable , closestack , frgt ) :
        while closestack and closestack[-1][0]<frgt :
            dummy,name = closestack.pop()
            if self.sql2xmltags :
                writable.write ('</%s>' % self.sql2xmltags(name))
            else : writable.write('</%s>' % name )

    def extract ( self , document , writable_ , comment=None ,
                  encoding='utf-8' ) :
        writable = EncodedFile(writable_,'utf-8',encoding)
        writable.write( '<?xml version="1.0" encoding="%s"?>' % encoding )
        if comment : writable.write ( '<!-- %d -->' % document )
        cr = self.cursor
        request = self.extraction_request ( document )
        cr.execute ( request )

        closestack = []

        row = cr.fetchone()
        if not row : return
        flft,frgt,name,value = row
        while 1 :
            self._stack_update ( writable , closestack , frgt )
            attrs = {name:value}
            oflft,ofrgt = flft,frgt

            while 1 :
                row = cr.fetchone()
                if not row : break
                flft,frgt,name,value = row
                if oflft != flft : break
                attrs[name]=value

            #print attrs
            tag_name = attrs.pop('','#UNDEFINED_TAG_NAME')
            if tag_name == '#TEXT' :
                writable.write(saxutils.escape(attrs.pop(
                    '#PCDATA','#UNDEFINED_PCDATA')))
            else :
                writable.write ( '<' )
                if self.sql2xmltags :
                    writable.write (self.sql2xmltags(tag_name))
                else :
                    writable.write ( tag_name )
                if self.sql2xmlattrs :
                    for k,v in attrs.iteritems() :
                        writable.write ( ' %s=%s' % (
                            self.sql2xmlattrs(k),
                            saxutils.quoteattr(v)) )
                else :
                    for k,v in attrs.iteritems() :
                        writable.write ( ' %s=%s' % (
                            k, saxutils.quoteattr(v)) )
                if oflft+1 == ofrgt : writable.write ( '/' )
                else : closestack.append ( (ofrgt,tag_name) )
                writable.write ( '>' )

            if not row : break
        self._stack_update ( writable , closestack , 256**9 )

    def extraction_request ( self , document ) :
        schema = self.schema
        union = [ """
        select 
        flft+lmap.offs as flft , 
        frgt+rmap.offs as frgt ,
        name , value
        from %(schema)s.document
        join %(schema)s.fragment_forest using(fragment) 
        join %(schema)s.map as lmap on (flft between lmap.lft and lmap.rgt)
        join %(schema)s.map as rmap on (frgt between rmap.lft and rmap.rgt)
        join %(schema)s.attr_list using (attr_list)
        join %(schema)s.attr_pair using (attr_pair)
        join %(schema)s.name using (name_id)
        join %(schema)s.value using (value_id)
        where lmap.map = document.map and rmap.map = document.map
        and document = %(document)d
        """ % vars() ]
        if self.tables is None : return union[0]+'\norder by 1\n'
        for t in self.tables.values() :
            cn = self._case_name ( t )
            cv = self._case_value ( t )
            cw = self._where ( t )
            ss = self._subselect ( t )
            table_name = t.table_name
            table_schema = t.schema
            union.append ( """
            select 
            flft+lmap.offs as flft , 
            frgt+rmap.offs as frgt ,
            %(cn)s as name ,
            %(cv)s as value 
            from %(schema)s.document
            join %(schema)s.fragment_forest using(fragment) 
            join %(schema)s.map as lmap on (flft between lmap.lft and lmap.rgt)
            join %(schema)s.map as rmap on (frgt between rmap.lft and rmap.rgt)
            join %(table_schema)s.%(table_name)s using (attr_list)
             , %(ss)s
             where lmap.map = document.map and rmap.map = document.map
             and document = %(document)d and %(cw)s
             """ % vars() )
        return '\nunion\n'.join(union)+'\norder by 1\n'
            
    def _case_name ( self , table ) :
        """
        Make case expression which convert table_ss.clmn number
        to sql field name
        """
        if len(table.fields)==1 : return "''"
        ss = table.table_name.strip()+'_ss'
        res = [ 'case %s.clmn' % ss ]
        for i,field in enumerate(table.fields[:-1]) :
            if field.sqlname() == 'attr_list' : # replace by table name sign
                res.append( "when %d then ''" % i )
            else :
                res.append( "when %d then '%s'" % (i,field.sqlname()) )
        if table.fields[-1].sqlname() == 'attr_list' : # replace by table name
            res.append( "else '' end" )
        else :
            res.append( "else '%s' end" % table.fields[-1].sqlname() )
        return ' '.join(res)

    def _case_value ( self , table ) :
        """
        Make case expression which convert table_ss.clmn number
        to sql field value
        """
        if len(table.fields)==1 : return "'%s'" % table.table_name
        ss = table.table_name.strip()+'_ss'
        res = [ 'case %s.clmn' % ss ]
        for i,field in enumerate(table.fields[:-1]) :
            if field.sqlname() == 'attr_list' : # replace by table name sign
                res.append( "when %d then '%s'" % (i,table.table_name) )
            else :
                res.append(
                    "when %d then %s.%s.%s::text" %
                    (i,table.schema,table.table_name,field.sqlname()) )
        if table.fields[-1].sqlname() == 'attr_list' : # replace by table name
            res.append( "else '%s' end" % table.table_name )
        else :
            res.append(
                "else %s.%s.%s::text end" %
                (table.schema,table.table_name,table.fields[-1].sqlname()) )
        return ' '.join(res)

    def _where ( self , table ) :
        """
        Make where condition to prevent using NULL values
        """
        if len(table.fields)==1 : return "TRUE" # table name is not NULL!
        ss = table.table_name.strip()+'_ss'
        res = [ ]
        for i,field in enumerate(table.fields) :
            if field.sqlname() == 'attr_list' : # replace by table name sign
                res.append( "(%s.clmn=%d)" % (ss,i) )
            else :
                res.append(
                    "(%s.clmn=%d and %s.%s.%s IS NOT NULL)" %
                    (ss,i,table.schema,table.table_name,field.sqlname()) )
        return '('+' or '.join(res)+')'

    def _subselect ( self , table ) :
        """
        Make enumeration request
        """
        ss = table.table_name.strip()+'_ss'
        res = ['select 0 as clmn' ]
        for i in range(1,len(table.fields)) :
            res.append ( 'select %d' % i )
        return '('+' union '.join(res)+') as '+ss
#-----------------------------------------------------------


if __name__ == '__main__' :
    import sys
    from main import main
    try: main()
    except IOError , KeyboardInterrupt : sys.exit(1) 
#===========================================================
